// Copyright 2022 jeff.li. and/or its affiliates.

#include <../src/codegen/codegen_c_host.h>
#include <tbir/ir/function.h>
#include <tbir/ir/module.h>
#include <tbir/ir/prim_builtin.h>
#include <tbir/ir/prim_expr.h>
#include <tbir/runtime/container.h>
#include <tbir/runtime/logging.h>
#include <tbir/runtime/registry.h>
#include <iostream>

#include <gtest/gtest.h>

namespace tbir::ir {
    using namespace runtime;

    TEST(IR, ForRange) {
        DataType f64_ty = DataType::Float(64);
        DataType i64_ty = DataType::Int(64);
        const auto *printer = ::tbir::runtime::FunctionRegistry::Get("ir.AsText");

        PrimVar arg_ib("ib", i64_ty);
        AllocaVarStmt alloca_result("result", PrimType(f64_ty), FloatImm(f64_ty, 0.0));
        AssignStmt set_res(alloca_result->var,
                           add(HLOCastPrim(f64_ty, alloca_result->var), FloatImm(f64_ty, 1.0)));

        For loop(PrimVar("i"), IntImm(i64_ty, 0), arg_ib, IntImm(i64_ty, 1), ForType::Serial, set_res);
        Array<Stmt> seq_stmt;
        seq_stmt.push_back(alloca_result);
        seq_stmt.push_back(loop);
        seq_stmt.push_back(ReturnStmt(alloca_result->var));
        SeqStmt body(seq_stmt);
        Array<BaseExpr> params{arg_ib};
        Function func(params, {}, body, PrimType(DataType::Float(64)), {});

        String ir_text = (*printer)({func}).As<String>();
        std::cout << ir_text << std::endl;
    }

}  // namespace tbir::ir
